package com.algorithm.cjm.csnote

import java.util.*


/**
 * 用PriorityQueue 实现小顶堆，加入k个元素，当元素大于k 过后，堆顶出去，最后堆顶的就是 第K 大的元素
 *
 * Created by jieming.chen on 2021/2/5
 */
fun findKthLargest(nums: IntArray, k: Int): Int {
    val pq = PriorityQueue<Int>() // 小顶堆

    for (element in nums) {
        pq.add(element)
        if (pq.size > k) // 维护堆的大小为 K
            pq.poll()
    }
    return pq.peek()
}

fun findKthLargest1(nums: IntArray, k: Int): Int {
    var k1 = nums.size - k
    var l = 0
    var h = nums.size - 1

    while (l < h) {
        var j = partition(nums, l, h)
        if (j == k1) {
            break
        } else if (j < k1) {
            l = j + 1
        } else {
            h = j - 1
        }
    }
    return nums[k1]
}

fun partition(nums: IntArray, l: Int, h: Int): Int {
    var i = l
    var j = h + 1

    while (true) {
        while (nums[++i] < nums[l] && i < h);
        while (nums[--j] > nums[l] && j > l);
        if (i >= j) {
            break
        }
        swap(nums, i, j)
    }
    swap(nums, l, j)
    return j
}

fun swap(a: IntArray, i: Int, j: Int) {
    val t: Int = a[i]
    a[i] = a.get(j)
    a[j] = t
}

fun main(args: Array<String>) {
    println(findKthLargest(intArrayOf(3, 2, 1, 5, 6, 4), 2))
    println(findKthLargest1(intArrayOf(3, 2, 1, 5, 6, 4), 2))
}